package com.example.base;

import cn.hutool.core.date.LocalDateTimeUtil;
import com.example.entity.Customer;
import com.example.entity.Order;
import com.example.entity.Payment;
import org.apache.flink.api.common.eventtime.TimestampAssigner;
import org.apache.flink.api.common.eventtime.TimestampAssignerSupplier;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.connector.datagen.source.DataGeneratorSource;
import org.apache.flink.connector.datagen.source.GeneratorFunction;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import java.time.Duration;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

/**
 * Author wangJinLong
 * Date 2025/8/15 09:51
 **/
public class FlinkStreamSource {

    private static final Long numbers_of_row = 100L;

    private static List<String> faceList = Arrays.asList("14901", "24902", "34903", "44903", "54903");
    private static List<String> siteList = Arrays.asList("05350101", "05350104", "05350122", "00004", "00002");
    private static List<String> nameList = Arrays.asList("tom", "jack", "rose", "ken", "tomas");

    public static DataStream<Order> getMapSource(StreamExecutionEnvironment env) {
        GeneratorFunction<Long, Order> generatorFunction = value -> {
            Order order = new Order();
            Random random = new Random();
            order.setOrderId(faceList.get(random.nextInt(faceList.size())));
            order.setAmount(random.nextInt(10) + 1);
            order.setTotalPrice(random.nextDouble(100));
            LocalDateTime now = LocalDateTime.now();
            order.setLocalOrderTime(LocalDateTimeUtil.format(now, "yyyy-MM-dd HH:mm:ss"));
            long epochMilli = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).toEpochMilli();

            int plus_minus = random.nextInt(2) % 2;
            long nextInt = (random.nextInt(10) + 1) * 1000;
            if (0 == plus_minus) {
                epochMilli = epochMilli - nextInt;
            }

            order.setOrderTime(epochMilli);
            return order;
        };

        DataGeneratorSource<Order> source = new DataGeneratorSource<>(
                generatorFunction,
                numbers_of_row,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Order.class)
        );

        DataStream<Order> mapSource = env.fromSource(
                source,
                WatermarkStrategy.<Order>forBoundedOutOfOrderness(Duration.ofSeconds(0))
                        .withTimestampAssigner(((element, recordTimestamp) -> element.getOrderTime()))
                , "map source"
        );

        return mapSource;
    }

    public static DataStream<Customer> getCustomerSource(StreamExecutionEnvironment env) {
        GeneratorFunction<Long, Customer> generatorFunction = value -> {
            Customer customer = new Customer();
            Random random = new Random();
            customer.setFeaturesId(faceList.get(random.nextInt(faceList.size())));
            return customer;
        };

        DataGeneratorSource<Customer> source = new DataGeneratorSource<>(
                generatorFunction,
                numbers_of_row,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Customer.class)
        );

        DataStream<Customer> mapSource = env.fromSource(
                source,
                WatermarkStrategy.<Customer>forBoundedOutOfOrderness(Duration.ofSeconds(0))
                        .withTimestampAssigner(((element, recordTimestamp) -> element.getCaptureTime()))
                , "Customer source"
        );

        return mapSource;
    }


    public static DataStream<Order> getOrderSource(StreamExecutionEnvironment env) {
        GeneratorFunction<Long, Order> generatorFunction = value -> {
            Order order = new Order();
            Random random = new Random();
            order.setOrderId(faceList.get(random.nextInt(faceList.size())));
            order.setAmount(random.nextInt(10) + 1);
            order.setTotalPrice(random.nextDouble(100));

            LocalDateTime now = LocalDateTime.now();
            order.setLocalOrderTime(LocalDateTimeUtil.format(now, "yyyy-MM-dd HH:mm:ss"));
            long epochMilli = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).toEpochMilli();

            int plus_minus = random.nextInt(2) % 2;
            long nextInt = (random.nextInt(3) + 1) * 1000;
            if (0 == plus_minus) {
                epochMilli = epochMilli - nextInt;
            }

            order.setOrderTime(epochMilli);
            return order;
        };

        DataGeneratorSource<Order> source = new DataGeneratorSource<>(
                generatorFunction,
                numbers_of_row,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Order.class)
        );

        DataStream<Order> mapSource = env.fromSource(
                source,
                WatermarkStrategy.<Order>forBoundedOutOfOrderness(Duration.ofSeconds(0))
                        .withTimestampAssigner(((element, recordTimestamp) -> element.getOrderTime()))
                , "map source"
        );

        return mapSource;
    }

    public static DataStream<Payment> getPaymentSource(StreamExecutionEnvironment env) {
        GeneratorFunction<Long, Payment> generatorFunction = value -> {
            Payment payment = new Payment();
            Random random = new Random();
            payment.setPayId(siteList.get(random.nextInt(5)));
            payment.setOrderId(faceList.get(random.nextInt(faceList.size())));
            payment.setPayName(nameList.get(random.nextInt(5)));
            payment.setPayAmount(random.nextInt(100));

            LocalDateTime now = LocalDateTime.now();
            payment.setLocalPayTime(LocalDateTimeUtil.format(now, "yyyy-MM-dd HH:mm:ss"));
            long epochMilli = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).toEpochMilli();

            int plus_minus = random.nextInt(2) % 2;
            long nextInt = (random.nextInt(5) + 1) * 1000;
            if (0 == plus_minus) {
                epochMilli = epochMilli - nextInt;
            }

            payment.setPayTime(epochMilli);
            return payment;
        };

        DataGeneratorSource<Payment> source = new DataGeneratorSource<>(
                generatorFunction,
                numbers_of_row,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Payment.class)
        );

        DataStreamSource<Payment> paymentSource = env.fromSource(
                source,
                WatermarkStrategy.<Payment>forBoundedOutOfOrderness(Duration.ofSeconds(0))
                        .withTimestampAssigner(((element, recordTimestamp) -> element.getPayTime())),
                "payment source"
        );

        return paymentSource;
    }
}
